import torch
import torch.nn as nn
from modules.vit_seg_modeling import VisionTransformer as transUnet
from modules.vit_seg_modeling import CONFIGS as CONFIGS_ViT_seg

# 封装为攻击模块
class TransUNet_Attack(nn.Module):
    def __init__(self):
        """
        """
        super(TransUNet_Attack, self).__init__()

        vit_name='R50-ViT-B_16'
        config_vit = CONFIGS_ViT_seg[vit_name]
        config_vit.n_classes = 3
        config_vit.n_skip = 3
        if vit_name.find('R50') != -1:
            config_vit.patches.grid = (int(128 / 16), int(128 / 16))
        self.unet = transUnet(config_vit, img_size=128, num_classes=config_vit.n_classes).cuda()
 
    def forward(self, x):

        noise_or_image = self.unet(x)

        return noise_or_image


